# @Time : 2021/8/3 17:25
# @Author : Li Kunlun
# @Description : 从零开始实现线性回归

from IPython import display
from matplotlib import pyplot as plt
from mxnet import autograd, nd
import random

# 1、生成数据集
num_inputs = 2
num_examples = 1000
true_w = [2, -3.4]
true_b = 4.2
features = nd.random.normal(scale=1, shape=(num_examples, num_inputs))
labels = true_w[0] * features[:, 0] + true_w[1] * features[:, 1] + true_b
labels += nd.random.normal(scale=0.01, shape=labels.shape)


def use_svg_display():
    # 用矢量图显示
    display.set_matplotlib_formats('svg')


def set_figsize(figsize=(3.5, 2.5)):
    use_svg_display()
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize


set_figsize()
# 通过生成第二个特征features[:, 1]和标签 labels 的散点图，可以更直观地观察两者间的线性关系。
plt.scatter(features[:, 1].asnumpy(), labels.asnumpy(), 1);  # 加分号只显示图,测试后差别不大
plt.show()


# 2、读取数据集
def data_iter(batch_size, features, labels):
    """
    :param
        :param batch_size: 小批量大小
        :param features: 特征x
        :param labels: 标签y
        :return: 小批量样本的特征和标签
    """
    num_examples = len(features)
    indices = list(range(num_examples))
    random.shuffle(indices)  # 样本的读取顺序是随机的
    for i in range(0, num_examples, batch_size):
        # 获取batch_size（10个数据）
        j = nd.array(indices[i: min(i + batch_size, num_examples)])
        yield features.take(j), labels.take(j)  # take函数根据索引返回对应元素


"""
1、测试
读取第一个小批量数据样本并打印。每个批量的特征形状为(10, 2)
2、输出结果
    [[-0.26484811  0.0712087 ]
 [ 0.95956105  0.14733684]
 [-1.09227788  0.80208206]
 [-0.2243541  -0.00421293]
 [-0.82268971 -1.45548177]
 [ 1.28527689 -1.82112122]
 [ 0.49716783 -0.36797494]
 [ 0.6378637  -0.69090891]
 [ 0.17943777  1.39908278]
 [-0.81124222  0.1232819 ]]
<NDArray 10x2 @cpu(0)> 
[  3.42115951   5.62397623  -0.72062272   3.75450468   7.51658964
  12.98480225   6.44769239   7.83661175  -0.20697819   2.16911864]
<NDArray 10 @cpu(0)>
"""
batch_size = 10
for X, y in data_iter(batch_size, features, labels):
    print(X, y)
    break

# 3、初始化模型函数 权重初始化成均值为0、标准差为0.01的正态随机数，偏差则初始化成0
w = nd.random.normal(scale=0.01, shape=(num_inputs, 1))
b = nd.zeros(shape=(1,))

# w-> (2, 1) b-> (1,)
# print("w->", w.shape, "b->", b.shape)

# 之后的模型训练中，需要对这些参数求梯度来迭代参数的值，因此需要创建它们的梯度
w.attach_grad()
b.attach_grad()


# 4、定义模型
def linreg(X, w, b):
    return nd.dot(X, w) + b


# 5、定义损失函数
def squared_loss(y_hat, y):
    return (y_hat - y.reshape(y_hat.shape)) ** 2 / 2


# 6、定义优化算法
def sgd(params, lr, batch_size):
    """
    :param
        :param params: [w, b]初始权重值
        :param lr: 学习率
        :param batch_size: 小批量大小
        :return: [w, b]
    """
    for param in params:
        param[:] = param - lr * param.grad / batch_size


# 7、训练模型
lr = 0.03
num_epochs = 3
net = linreg
loss = squared_loss

for epoch in range(num_epochs):  # 训练模型一共需要num_epochs个迭代周期
    """
     1、在每一个迭代周期中，会使用训练数据集中所有样本一次（假设样本数能够被批量大小整除）。
     2、X和y分别是小批量样本的特征和标签
    """
    for X, y in data_iter(batch_size, features, labels):
        with autograd.record():
            l = loss(net(X, w, b), y)  # l是有关小批量X和y的损失
        l.backward()  # 小批量的损失对模型参数求梯度
        sgd([w, b], lr, batch_size)  # 使用小批量随机梯度下降迭代模型参数
    train_l = loss(net(features, w, b), labels)

    # epoch 1, loss 0.035158
    # epoch 2, loss 0.000131
    # epoch 3, loss 0.000048
    print('epoch %d, loss %f' % (epoch + 1, train_l.mean().asnumpy()))

# 计算结果和原始结果相差不大
print("true_w->", true_w, "w->", w)
print("true_b->", true_b, "b->", b)
